import time
import torch
import argparse
from dataset.dataset2 import GraphDataset, read_data_trackA, read_data_trackB
import numpy as np

class LpLoss(object):
    def __init__(self, d=2, p=2, size_average=True, reduction=True):
        super(LpLoss, self).__init__()

        # Dimension and Lp-norm type are postive
        assert d > 0 and p > 0

        self.d = d
        self.p = p
        self.reduction = reduction
        self.size_average = size_average

    def abs(self, x, y):
        num_examples = x.size()[0]

        # Assume uniform mesh
        h = 1.0 / (x.size()[1] - 1.0)

        all_norms = (h ** (self.d / self.p)) * torch.norm(
            x.view(num_examples, -1) - y.view(num_examples, -1), self.p, 1
        )

        if self.reduction:
            if self.size_average:
                return torch.mean(all_norms)
            else:
                return torch.sum(all_norms)

        return all_norms

    def rel(self, x, y):
        num_examples = x.size()[0]

        diff_norms = torch.norm(
            x.reshape(num_examples, -1) - y.reshape(num_examples, -1), self.p, 1
        )
        y_norms = torch.norm(y.reshape(num_examples, -1), self.p, 1)

        if self.reduction:
            if self.size_average:
                return torch.mean(diff_norms / y_norms)
            else:
                return torch.sum(diff_norms / y_norms)

        return diff_norms / y_norms

    def __call__(self, x, y):
        return self.rel(x, y)


parser = argparse.ArgumentParser()
parser.add_argument('--data_dir', default='dataset/Track_B')
parser.add_argument('--info_dir', default='dataset/Track_B')
parser.add_argument('--fold_id', default=0, type=int)
parser.add_argument('--gpu', default=0, type=int)
parser.add_argument('--val_iter', default=10, type=int)
parser.add_argument('--cfd_model', default='TransAFNO')
parser.add_argument('--cfd_mesh', default=True)
parser.add_argument('--r', default=0.2, type=float)
parser.add_argument('--weight', default=0.5, type=float)
parser.add_argument('--lr', default=0.001, type=float)
parser.add_argument('--batch_size', default=1, type=float)
parser.add_argument('--nb_epochs', default=1000, type=float)
parser.add_argument('--preprocessed', default=0, type=int)
args = parser.parse_args()
print(args)

if 'Track_A' in args.data_dir:
    train_data, val_data, coef_norm, val_index = read_data_trackA(args, norm=True)
else:
    train_data, val_data, coef_norm, val_index = read_data_trackB(args, norm=True)

n_gpu = torch.cuda.device_count()
use_cuda = 0 <= args.gpu < n_gpu and torch.cuda.is_available()
device = torch.device(f'cuda:{args.gpu}' if use_cuda else 'cpu')


from models.TransAFNO import Model
model = Model(n_hidden=128, n_layers=4, space_dim=3,
              fun_dim=0,
              n_head=8,
              mlp_ratio=2, out_dim=1,
              slice_num=16,
              unified_pos=0).cuda()

path = 'metrics/model_0.15398150891541396.pth'
model = torch.load(path).to(device)

with torch.no_grad():
    model.eval()

    criterion_func = LpLoss(size_average=True)
    l2errs_press = []
    times = []
    index = 0
    for cfd_data in val_data:
        cfd_data = cfd_data.to(device)
        tic = time.time()
        out = model((cfd_data, cfd_data))
        toc = time.time()
        targets = cfd_data.y

        if coef_norm is not None:
            mean = torch.tensor(coef_norm[2]).to(device)
            std = torch.tensor(coef_norm[3]).to(device)
            pred_press = out * std[-1] + mean[-1]
            gt_press = targets * std[-1] + mean[-1]
            l2_loss = criterion_func(pred_press.unsqueeze(0), gt_press.unsqueeze(0).unsqueeze(-1))
            print('l2 loss:', l2_loss.item())
            l2errs_press.append(l2_loss.detach().cpu().numpy())
            np.save(f"./output/answer/press_{val_index[index]}.npy", pred_press.detach().cpu().numpy().squeeze())
            index = index + 1

        times.append(toc - tic)

    print('time:', np.mean(times))
    print('l2errs_press:', np.mean(l2errs_press))

